# 检测一个基因条码 是否为consensus by length
# 如果是 输出 true

from Bio import SeqIO
from Bio.Seq import translate
import re

def check_consensus_by_length(seq,l=658):
    """这个基因条码的长度是否符合 制定长度 不符合 输出 false"""

    if len(seq) != l:
        return False
    # 检查是否可以被翻译
    return True

def check_consensus_by_translation(seq,table=5):
    """这个条码是否可以被翻译 不符合 False 如果符合 输出 pep"""
    n = False
    for frame_start in range(3):
        frame = translate(seq[frame_start:frame_start+(len(seq)-frame_start)//3*3],table=table)
        #print(frame)
        if frame.find('*')==-1:
            n = frame
            break
    return n


def check_frame_translation(seq,table=5):
    """查看能翻译的阅读框"""
    n = False
    for frame_start in range(3):
        frame = translate(seq[frame_start:frame_start+(len(seq)-frame_start)//3*3],table=table)
        #print(frame)
        if frame.find('*')==-1:
            n = frame_start
            break
    return n

def check_N(seq):
    if seq.find('N')!=-1 or seq.find('n')!=-1:
        return False
    return True


# infile = 'consensus_reference_11.fasta'

# for i in SeqIO.parse(infile,'fasta'):
#     k = check_consensus_by_translation(i.seq)
#     print(k)



def aliToGapSeq(seq_lista,t = 0.5):

    '''将多序列比对 转化成一条 能反应其中体gap的 序列'''
    return_seq = ''
    for i in range(len(seq_lista[0])):
        site_lista = [seq[i] for seq in seq_lista]
        if site_lista.count('-')/len(seq_lista)>t:
            return_seq+='-'
        else:
            return_seq += '.'
    return return_seq


def get_gap_codon_block(seq,frame):
    lista = []

    hy = re.compile('[-]+')
    for i in hy.finditer(seq):
        start_point = i.start(0)
        end_point = i.end(0)
        if (start_point-end_point) %3 == 0:
            #l = (start-point) % 3
            frame_shift_start = (start_point-frame) % 3
            frame_shift_end = 3 - frame_shift_start 

            lista.append([start_point - frame_shift_start,end_point +frame_shift_end, start_point,end_point])

    return lista


def mafft2GapSeq(mafft_file, table=5):
    '''读取mafft 文件 假定 第一条序列为 目标序列 后面的所有序列为参考序列
        这个函数要判断的是 这个target序列 和reference序列 他们是什么关系
        [fixed_seq,[errors]]

    '''

    ref_seq = []

    k = SeqIO.parse(mafft_file,'fasta')
    try:
        first_seq = next(k).seq
    except:
        return False


    result_lista =[str(first_seq).upper(),[]]
    
    for i in k:
        ref_seq.append(i.seq)
    #print(ref_seq)
    
    if ref_seq   == []:
        return False

    # 将ref seq 的结果转化为 gapseq
    ref_gapseq = aliToGapSeq(ref_seq,0.5)


    t = ref_gapseq.strip('-')
    no_bar_range = [ref_gapseq.find(t),ref_gapseq.find(t)+len(t)]

    if no_bar_range[0] != 0:
        result_lista[1].append('side mismatch 5p size='+str(no_bar_range[0]))
        
    if no_bar_range[1] != len(ref_gapseq):
        result_lista[1].append('side mismatch 3p size='+str(len(result_lista[0])-no_bar_range[1]))

    result_lista[0] = result_lista[0][no_bar_range[0]:no_bar_range[1]]
    ref_gapseq = ref_gapseq[no_bar_range[0]:no_bar_range[1]]
    
    # 查看内部 indel error


    if result_lista[0].find('-')!=-1 or ref_gapseq.find('-')!=-1:

        frame = check_frame_translation(str(ref_seq[0]).replace('-',''))
        # print(ref_gapseq)
        if frame:
            # mask gap on the ref
            l = get_gap_codon_block(ref_gapseq,frame)
            #print(result_lista[0])
            #print(ref_gapseq)
            for j in l:
                #print(ref_gapseq[j[0]:j[1]])
                #print(mafft_file)

                if result_lista[0][j[2]:j[3]].find('-')==-1 and translate(result_lista[0][j[2]:j[3]]).find('*')==-1:
                    # 如果是 全部mask
                    ref_gapseq = ref_gapseq[:j[0]] + '.'*(j[1]-j[0])+ ref_gapseq[j[1]:]
                    result_lista[1].append('insertion codon at %s to %s'%(j[2],j[3]))
            # 删除缺失 这里改为 $ 最后在清除 目的是为了保证 后面的 deletion 和insertion error 的位置 都对应到比对上


            l = get_gap_codon_block(result_lista[0],frame)

            for j in l:
                # 这里需要保证 gap 不同时出现在 原序列中
                result_lista[0] = result_lista[0][:j[2]] + '$'*(j[3]-j[2])+ result_lista[0][j[3]:]
                result_lista[1].append('deletion codon at %s to %s'%(j[2],j[3]))

        # print(ref_gapseq)
        

        target_seq = ''

        n = 0
        while n < len(result_lista[0]):
                
            if result_lista[0][n] =='-' and ref_gapseq[n]=='-':
                pass
            elif result_lista[0][n] =='$' and ref_gapseq[n]=='-':
                pass
            elif result_lista[0][n]=='-' and ref_gapseq[n]!='-':
                target_seq += 'N'
                result_lista[1].append('deletion error at '+str(n+no_bar_range[0]))
            elif result_lista[0][n]!='-' and ref_gapseq[n]=='-':
                result_lista[1].append('insertion error at '+str(n+no_bar_range[0]))
            else:
                target_seq += result_lista[0][n]
            
            n+=1

        result_lista[0] = target_seq.replace('$','')

    return result_lista


#print(mafft2GapSeq('/media/good/myData2021/barcode_work/sample1/ONT/test123/COI_F_8G_R_1E/consensus_reference_2/mafft_out.fasta',5))
